#include <stdlib.h>
#include <stdio.h>
#include <stdbool.h>
#include <signal.h>
#include <unistd.h>
#include <setjmp.h>
#include <string.h>
#include <assert.h>

#include "list.h"
#include "user_mutex_inner.h"
#include "user_common_inner.h"
#include "user_threads_inner.h"
#include "user_threads.h"

/* the maximum number of threads allowed */
#define NR_THREADS    200

#define LOCKED (lock_count > 0)

/* thread's state */
enum thread_state {
	UTHREAD_NOT_INITIALIZED = -1,
	UTHREAD_RUNNING = 0,
	UTHREAD_SUSPENDED = 1,
	UTHREAD_ZOMBIE = 2,
};

/* context switch reason */
enum cs_reason {
	CS_SCHED,
	CS_WAIT,
	CS_YIELD,
	CS_KILL,
};

#define WAIT_NONE -1

#define curr get_current()

/* thread descriptor */
typedef struct uthread_t {
    int state;
    thread_id id;

	sigjmp_buf jbuf; /* used by sigjmplong */

    void *stack;

	UTHREAD_START_ROUTINE start;
	void* arg;

	int ret_val;

	thread_id wait_for;
	thread_id waited_by;
	int join_ret_val;

	list locks; /* mutex locks list */

	list *current_queue; /* the current run/wait queue */
} uthread;

static bool initialized = false;

static uthread uthreads[NR_THREADS];

static list rq;

static int lock_count = 0;

static __useconds_t time_left = TIME_SLICE;

static uthread* get_current();
static uthread* it_to_uthread(iterator *it);
static void reset_thread(uthread *t);
static void run_thread();
static void kill_thread(uthread *t, int ret_val);
static void schedule(int);
static void context_switch(uthread *from, uthread *to, int reason);
static void context_switch_next(uthread *from, int reason);
static void resume_thread(thread_id tid);
static void remove_from_rq(uthread *t, list *wait_queue);
static void clear_zombie(uthread *t);
static void register_sig_handler();
static iterator get_iterator_in_current_queue(uthread *t);
static uthread* find_next_running(uthread *t);

void uthread_init()
{
	int i;
	uthread *t;
	int lrc;
	int ret = MUTEX_FAILURE;

    if (initialized)
    	return;

    lock_count = 1;

    for (i = 0; i < NR_THREADS; ++i) {
    	reset_thread(uthreads + i);
    }

    lrc = list_init(&rq);
    assert(lrc != LIST_ERROR);

    t = uthreads;
    t->id = 0;
    t->state = UTHREAD_SUSPENDED;

    lrc = list_init(&(t->locks));
    assert(lrc != LIST_ERROR);

    t->current_queue = NULL;
    resume_thread(t->id);

    ret = mutexes_init();
    assert(MUTEX_SUCCESS == ret);

    initialized = true;

    register_sig_handler();
    uthread_unlock();

    assert(!LOCKED);
}

thread_id uthread_create(UTHREAD_START_ROUTINE start_routine, void *arg)
{
    int ret = UTHREAD_FAIL;
    uthread *t = NULL;
    int lrc;

    uthread_init();

    uthread_lock();

    DPRINT("uthread_create()");

    /* find the next currently free id */
    int next_tid = 0;
    while((next_tid < NR_THREADS) && (uthreads[next_tid].state > UTHREAD_NOT_INITIALIZED))
		++next_tid;

    if (next_tid >= NR_THREADS)
        goto out;

    /* initialize new thread descriptor */
    t = uthreads + next_tid;

    t->stack = malloc(UTHREAD_DEF_STACK_SIZE);
    if (!(t->stack))
    	goto out;

    t->id = next_tid;
    t->start = start_routine;
	t->arg = arg;

    lrc = list_init(&t->locks);
    assert(lrc != LIST_ERROR);

    /* set the thread jump parameters */
    sigsetjmp(t->jbuf, 1);
    t->jbuf->__jmpbuf[JB_PC] = (int) run_thread;
    t->jbuf->__jmpbuf[JB_SP] = (int) ((char*) t->stack + UTHREAD_DEF_STACK_SIZE);
    t->jbuf->__jmpbuf[JB_BP] = (int) ((char*) t->stack + UTHREAD_DEF_STACK_SIZE);

    /* make the new thread run last in the RR */
    t->state = UTHREAD_SUSPENDED;
    t->current_queue = NULL;
    resume_thread(t->id);

    ret = t->id;

out:
	uthread_unlock();
	assert(!LOCKED);
    return ret;
}

void uthread_exit(int retval)
{
	uthread_init();

	uthread_lock();

	DPRINT("uthread_exit() on %d", curr->id);

	kill_thread(curr, retval);
	assert(0); /* we should never get here */
}

thread_id uthread_self()
{
	thread_id tid;

	uthread_init();

	uthread_lock();
    tid = curr->id;
    uthread_unlock();
    return tid;
}

int uthread_join(thread_id tid, int *thread_ret_val)
{
	uthread *t;
	uthread *prev;
	int ret = UTHREAD_FAIL;

	uthread_init();

	uthread_lock();

	DPRINT("uthread_join(%d) from tid %d", tid, curr->id);

	if (tid == curr->id)
		goto out;

	ret = UTHREAD_INVALID;
	if (tid < 0 || tid > NR_THREADS)
		goto out;

	t = uthreads + tid;
	if (t->state == UTHREAD_NOT_INITIALIZED)
		goto out;

	ret = UTHREAD_FAIL;
	if (t->waited_by != WAIT_NONE)
		goto out;

	ret = UTHREAD_SUCCESS;

	if (t->state != UTHREAD_ZOMBIE) {
		assert(curr->wait_for == WAIT_NONE);
		curr->wait_for = tid;
		t->waited_by = curr->id;

		if (list_get_size(&rq) <= 1) {
			/* "occured" misspelled as in the FAQ... */
			DPRINT("!!! DEADLOCK !!! - exiting");
			printf("deadlock occured");
			exit(0);
		}

		/* make it wait */
		prev = curr;
		remove_from_rq(prev, NULL);
		context_switch_next(prev, CS_WAIT);

		uthread_lock(); /* re-lock after returning from wait */

		/* at this point, t might be not initialized or a new created thread */

		assert(prev == curr);
		assert(curr->wait_for == WAIT_NONE);

		if (thread_ret_val)
			*thread_ret_val = curr->join_ret_val;
	}
	else { /* UTHREAD_ZOMBIE */
		if (thread_ret_val)
			*thread_ret_val = t->ret_val;

		clear_zombie(t);
	}

out:
	uthread_unlock();
	assert(!LOCKED);
    return ret;
}

int uthread_cancel(thread_id tid)
{
	uthread *t;
	int ret = UTHREAD_INVALID;

	uthread_init();

	uthread_lock();

	DPRINT("uthread_cancel(%d)", tid);

	if (tid < 0 || tid >= NR_THREADS)
		goto out;

	t = uthreads + tid;
	if (t->state == UTHREAD_NOT_INITIALIZED) /* not a valid thread */
		goto out;

	ret = UTHREAD_FAIL;
	if (t->state == UTHREAD_ZOMBIE) /* it's already canceled */
		goto out;

	kill_thread(t, UTHREAD_CANCELLED);
	ret = UTHREAD_SUCCESS;

out:
	uthread_unlock();
	assert(!LOCKED);
	return ret;
}

int uthread_yield()
{
	int ret;
	int lrc;

	uthread_init();

	uthread_lock();

	DPRINT("uthread_yield() on %d", curr->id);

    if (list_get_size(&rq) > 1) {
    	uthread *prev = curr;
    	uthread *next = find_next_running(curr);
    	iterator next_it = get_iterator_in_current_queue(next);

    	assert(next != curr);

    	lrc = list_remove(&rq, &next_it);
    	assert(lrc == LIST_SUCCESS);

		lrc = list_remove_head(&rq);
		assert(lrc == LIST_SUCCESS);

		lrc = list_insert_tail(&rq, prev);
		assert(lrc == LIST_SUCCESS);

    	lrc = list_insert_head(&rq, next);
		assert(lrc == LIST_SUCCESS);

		/* make the next thread run */
		context_switch_next(prev, CS_YIELD);
		ret = UTHREAD_SUCCESS;
    }
    else {
    	time_left = TIME_SLICE; /* prevent the next schedule() since there is nothing to schedule */
    	uthread_unlock();
    	ret = UTHREAD_FAIL; /* no context switch occurred */
    }

    assert(!LOCKED);
    return ret;
}

int uthread_suspend_self(list *wait_queue)
{
	uthread *t = curr;

	assert(LOCKED);
	remove_from_rq(t, wait_queue);
	context_switch_next(t, CS_WAIT);

	return UTHREAD_SUCCESS;
}

int uthread_resume_from_wait_queue(list *wait_queue, thread_id *resumed_tid)
{
	uthread *t;
    iterator it;
    int lrc;

    assert(LOCKED);

	lrc = list_get_head(wait_queue, &it);
    assert(lrc == LIST_SUCCESS);

    t = it_to_uthread(&it);
    assert (t->current_queue == wait_queue);

    lrc = list_remove(wait_queue, &it);
    assert(lrc != LIST_ERROR);

	t->current_queue = NULL;
    resume_thread(t->id);

    *resumed_tid = t->id;

	return UTHREAD_SUCCESS;
}

static uthread* get_current()
{
	iterator it;

	assert(LOCKED);
	int lrc = list_get_head(&rq, &it);
	assert(lrc == LIST_SUCCESS);
	return lrc == LIST_SUCCESS ? it_to_uthread(&it) : NULL;
}

static uthread* it_to_uthread(iterator *it)
{
	uthread *t;

	assert(LOCKED);
	int lrc = iterator_get_value(it, (list_element_t *) &t);
	assert(lrc == LIST_SUCCESS);
	return lrc == LIST_SUCCESS ? t : NULL;
}

static void resume_thread(thread_id tid)
{
	uthread *t = uthreads + tid;
	int lrc;

	assert(LOCKED);
	assert(t->state == UTHREAD_SUSPENDED); /* assert we came from suspension / creation */

	DPRINT("adding %d to rq", tid);

	assert(t->current_queue == NULL);
    t->current_queue = &rq;
    t->state = UTHREAD_RUNNING;

	lrc = list_insert_tail(&rq, (list_element_t) t);
    assert(lrc != LIST_ERROR);
}

static uthread* find_next_running(uthread *t)
{
	assert(LOCKED);
	assert(t != NULL);
	assert(t->state == UTHREAD_RUNNING);

	int next_tid = t->id + 1;

	/* look forward */
	while ((next_tid < NR_THREADS) && (uthreads[next_tid].state != UTHREAD_RUNNING))
		next_tid++;

	/* if did not found, start looking from the start */
	if (next_tid >= NR_THREADS)
		next_tid = 0;

	while ((next_tid < t->id) && (uthreads[next_tid].state != UTHREAD_RUNNING))
		next_tid++;

	return uthreads + next_tid;
}

static iterator get_iterator_in_current_queue(uthread *t)
{
	iterator it;

	assert(LOCKED);
	assert(t != NULL);
	assert(t->current_queue != NULL);
	assert(t->state != UTHREAD_NOT_INITIALIZED);

	LIST_FOREACH(t->current_queue, &it) {
		if (it_to_uthread(&it) == t)
			break;
	}

	assert(!(list_is_end(t->current_queue, &it)));
	return it;
}

static void reset_thread(uthread *t)
{
	assert(LOCKED);

	memset(t, 0, sizeof(uthread));
	t->state = UTHREAD_NOT_INITIALIZED;
	t->wait_for = WAIT_NONE;
	t->waited_by = WAIT_NONE;
}

static void run_thread()
{
	UTHREAD_START_ROUTINE start;
	void *arg;

	uthread_lock();

	start = curr->start;
	arg = curr->arg;

	uthread_unlock();

	uthread_exit(start(arg));
}

static void kill_thread(uthread *t, int ret_val)
{
	iterator it;
	int lrc;
	bool resched = false;

	assert(LOCKED);
	assert(t->state == UTHREAD_RUNNING || t->state == UTHREAD_SUSPENDED);

	t->ret_val = ret_val;

	if (t->state == UTHREAD_RUNNING) {
		resched = t == curr;
		remove_from_rq(t, NULL);
	}
	else if (t->current_queue != NULL) {
		/* if thread is waiting, remove it from wait queue */
		iterator it = get_iterator_in_current_queue(t);
		lrc = list_remove(t->current_queue, &it);
		assert(lrc == LIST_SUCCESS);
	}

	DPRINT("killing thread %d", t->id);

	t->state = UTHREAD_ZOMBIE;

	/* destroy the locks */
	LIST_FOREACH(&t->locks, &it) {
		uthread_mutex_t m;
		lrc = iterator_get_value(&it, (list_element_t *) &m);
		assert(lrc == LIST_SUCCESS);
		mutex_set_zombie(m);
	}

	lrc = list_destroy(&t->locks);
	assert(lrc != LIST_ERROR);

	/* resume joined thread */
	if (t->waited_by != WAIT_NONE) {
		uthread *jt = uthreads + t->waited_by;

		assert(jt->state == UTHREAD_SUSPENDED || jt->state == UTHREAD_ZOMBIE);
		assert(jt->wait_for == t->id);

		if (jt->state == UTHREAD_SUSPENDED) {
			DPRINT("resuming joined %d", jt->id);
			resume_thread(jt->id);
		}

		jt->join_ret_val = t->ret_val;
		jt->wait_for = WAIT_NONE;

		clear_zombie(t);
	}

	/* if killed curr, make the next thread run */
	if (resched) {
		if (list_get_size(&rq) > 0)
			context_switch_next(t, CS_KILL);
		else {
			/* no more joining threads should exist */
			DPRINT("no threads left - exiting");
			exit(ret_val);
		}
	}
}

static void schedule(int n)
{
	uthread_lock();

	assert(lock_count == 1);
	assert(curr != NULL);
	assert(list_get_size(&rq) > 0);

	DPRINT("schedule() for tid=%d", curr->id);

	if (list_get_size(&rq) > 1) {
		iterator next_it;
		uthread* t = curr;
		int lrc;

		lrc = list_remove_head(&rq);
		assert(lrc == LIST_SUCCESS);

		list_insert_tail(&rq, t);
		assert(lrc == LIST_SUCCESS);

		list_get_head(&rq, &next_it);
		assert(lrc == LIST_SUCCESS);

		context_switch(t, it_to_uthread(&next_it), CS_SCHED);
	}
	else {
		uthread_unlock();
	}

	assert(!LOCKED);
}

#ifdef ENABLE_DPRINT
static char* get_cs_reason_string(int reason)
{
	switch (reason) {
	case CS_SCHED:
		return "CS_SCHED";
	case CS_WAIT:
		return "CS_WAIT";
	case CS_YIELD:
		return "CS_YIELD";
	case CS_KILL:
		return "CS_KILL";
	}

	return "?????";
}
#endif

static void context_switch(uthread *from, uthread *to, int reason)
{
	assert(lock_count == 1);
	assert(from != NULL);
	assert(to != NULL);
	assert(from != to);
	assert(to->state == UTHREAD_RUNNING);

	DPRINT("switching from %d to %d (%s)", from->id, to->id, get_cs_reason_string(reason));

	if (sigsetjmp(from->jbuf, 1) == 0) {
		time_left = TIME_SLICE;
		uthread_unlock();
		assert(!LOCKED);
		siglongjmp(to->jbuf, 1);
	}

	assert(!LOCKED);
}

static void context_switch_next(uthread *from, int reason)
{
	iterator next_it;
	int lrc;

	assert(LOCKED);
	lrc = list_get_head(&rq, &next_it);
	assert(lrc == LIST_SUCCESS);
	assert(from != it_to_uthread(&next_it));

	context_switch(from, it_to_uthread(&next_it), reason);
}

static void remove_from_rq(uthread *t, list *wait_queue)
{
	iterator it;
	int lrc;

	assert(LOCKED);
	assert(t->state == UTHREAD_RUNNING);
	assert(t->current_queue == &rq);

	DPRINT("removing %d from rq", t->id);

	it = get_iterator_in_current_queue(t);
	lrc = list_remove(&rq, &it);
	assert(lrc == LIST_SUCCESS);

	t->state = UTHREAD_SUSPENDED;

	/* add to new suspend queue if given */
	t->current_queue = wait_queue;
	if (t->current_queue) {
		lrc = list_insert_tail(wait_queue, t);
	    assert(lrc == LIST_SUCCESS);
	}
}

static void clear_zombie(uthread *t)
{
	assert(LOCKED);
	assert(t->state == UTHREAD_ZOMBIE);

	/* notify the joined thread we don't wait for it anymore */
	if (t->wait_for != WAIT_NONE) {
		uthread *jt = uthreads + t->wait_for;
		assert(jt->waited_by == t->id);
		jt->waited_by = WAIT_NONE;
	}

	free(t->stack);
	reset_thread(t);
}

static void register_sig_handler()
{
	struct sigaction sa;

	memset(&sa, 0, sizeof(sa));
	sa.sa_handler = schedule;
	sigaction(SIGALRM, &sa, NULL);
}

void uthread_lock()
{
	if (lock_count++ == 0) {
		time_left = ualarm(0, 0);

		if (!time_left)
			time_left = TIME_SLICE;
	}
}

void uthread_unlock()
{
	if (--lock_count == 0)
		ualarm(time_left, 0);
}

list* uthread_get_locks(thread_id tid)
{
	assert(tid >= 0);
	assert(tid < NR_THREADS);

	uthread *t = uthreads + tid;
	return &(t->locks);
}
